#include <ntddk.h>
#include <wdm.h>

typedef unsigned int UINT;
typedef char * PCHAR;

BOOLEAN IsStringTerminated(PCHAR pString, UINT uiLength);
NTSTATUS UnSupportedFunction(PDEVICE_OBJECT DeviceObject, PIRP Irp);
VOID Unload(PDRIVER_OBJECT  DriverObject);
NTSTATUS DriverWrite(PDEVICE_OBJECT DeviceObject, PIRP Irp);
NTSTATUS DriverRead(PDEVICE_OBJECT DeviceObject, PIRP Irp);
NTSTATUS Close(PDEVICE_OBJECT DeviceObject, PIRP Irp);



NTSTATUS DriverEntry(PDRIVER_OBJECT  pDriverObject, PUNICODE_STRING  pRegistryPath)
{
    NTSTATUS NtStatus = STATUS_SUCCESS;
    UINT uiIndex = 0;
    PDEVICE_OBJECT pDeviceObject = NULL;
    UNICODE_STRING usDriverName, usDosDeviceName;

    DbgPrint("DriverEntry Called \r\n");

    RtlInitUnicodeString(&usDriverName, L"\\Device\\TestDriver");
    RtlInitUnicodeString(&usDosDeviceName, L"\\DosDevices\\TestDriver"); 

    NtStatus = IoCreateDevice(pDriverObject, 0, &usDriverName, FILE_DEVICE_UNKNOWN, FILE_DEVICE_SECURE_OPEN, FALSE, &pDeviceObject);

    if (NtStatus == STATUS_SUCCESS)
    {
        for(uiIndex = 0; uiIndex < IRP_MJ_MAXIMUM_FUNCTION; uiIndex++)
             pDriverObject->MajorFunction[uiIndex] = UnSupportedFunction;
    	
    	//DbgPrint("max: %d", IRP_MJ_MAXIMUM_FUNCTION);

        pDriverObject->MajorFunction[IRP_MJ_READ]              = DriverRead;
        pDriverObject->MajorFunction[IRP_MJ_WRITE]             = DriverWrite;
    	pDriverObject->MajorFunction[IRP_MJ_CLOSE]             = Close;

        pDriverObject->DriverUnload = Unload; 
        pDeviceObject->Flags |= DO_DIRECT_IO;
        pDeviceObject->Flags &= (~DO_DEVICE_INITIALIZING);

        IoCreateSymbolicLink(&usDosDeviceName, &usDriverName);
    }
    return NtStatus;
}


NTSTATUS Close(PDEVICE_OBJECT DeviceObject, PIRP Irp)
{
    NTSTATUS NtStatus = STATUS_SUCCESS;
    DbgPrint("Close Called \r\n");
    return NtStatus;
}



VOID Unload(PDRIVER_OBJECT  DriverObject)
{    
    UNICODE_STRING usDosDeviceName;
    DbgPrint("Unload Called \r\n");
    RtlInitUnicodeString(&usDosDeviceName, L"\\DosDevices\\TestDriver");
    IoDeleteSymbolicLink(&usDosDeviceName);
    IoDeleteDevice(DriverObject->DeviceObject);
}



NTSTATUS DriverRead(PDEVICE_OBJECT DeviceObject, PIRP Irp)
{
    NTSTATUS NtStatus = STATUS_BUFFER_TOO_SMALL;
    PIO_STACK_LOCATION pIoStackIrp = NULL;
    PCHAR pReturnData = "read performed";
    UINT dwDataSize = sizeof("read performed");
    UINT dwDataRead = 0;
    PCHAR pReadDataBuffer;

    DbgPrint("DriverRead Called\n");
    
    pIoStackIrp = IoGetCurrentIrpStackLocation(Irp);
    
    if(pIoStackIrp && Irp->MdlAddress)
    {
        pReadDataBuffer = MmGetSystemAddressForMdlSafe(Irp->MdlAddress, NormalPagePriority);
    
        if(pReadDataBuffer && pIoStackIrp->Parameters.Read.Length >= dwDataSize)
        {                             
            RtlCopyMemory(pReadDataBuffer, pReturnData, dwDataSize);
            DbgPrint("copied!!!");
            dwDataRead = dwDataSize;
            NtStatus = STATUS_SUCCESS;
        }
    }
    
    Irp->IoStatus.Status = NtStatus;
    Irp->IoStatus.Information = dwDataRead;

    IoCompleteRequest(Irp, IO_NO_INCREMENT);

    return NtStatus;
}


NTSTATUS DriverWrite(PDEVICE_OBJECT DeviceObject, PIRP Irp)
{
    NTSTATUS NtStatus = STATUS_SUCCESS;
    PIO_STACK_LOCATION pIoStackIrp = NULL;
    PCHAR pWriteDataBuffer;
    DbgPrint("DriverWrite Called\n");
    pIoStackIrp = IoGetCurrentIrpStackLocation(Irp);
    if (pIoStackIrp)
    {
        pWriteDataBuffer = MmGetSystemAddressForMdlSafe(Irp->MdlAddress, NormalPagePriority);
        if (pWriteDataBuffer)
        {
            if (IsStringTerminated(pWriteDataBuffer, pIoStackIrp->Parameters.Write.Length))
            {
                DbgPrint("writing: ");
                DbgPrint(pWriteDataBuffer);
            }
        }
    }
    return NtStatus;
}



NTSTATUS UnSupportedFunction(PDEVICE_OBJECT DeviceObject, PIRP Irp)
{
    NTSTATUS NtStatus = STATUS_NOT_SUPPORTED;
    DbgPrint("UnSupportedFunction Called \r\n");
    return NtStatus;
}


BOOLEAN IsStringTerminated(PCHAR pString, UINT uiLength)
{
    BOOLEAN bStringIsTerminated = FALSE;
    UINT uiIndex = 0;

    while(uiIndex < uiLength && bStringIsTerminated == FALSE)
    {
        if(pString[uiIndex] == '\0')
        {
            bStringIsTerminated = TRUE;
        }
        else
        {
           uiIndex++;
        }
    }

    return bStringIsTerminated;
}